from solvers.solverbase import SolverBase
from solvers.tg.tgbase import TGBase
from dolfin import *

class LaxWendroffTG(SolverBase,TGBase):
    'Lax-Wendroff Taylor-Galerkin method.'
    def __init__(self,p=1,type=0,bcStrong=0):
        SolverBase.__init__(self,p)
        TGBase.__init__(self,type,bcStrong)

    def solve(self,problem):
        mesh = problem.domain

        V = FunctionSpace(mesh,'CG',1)
        VV = VectorFunctionSpace(mesh,'CG',1)

        phi_ = interpolate(problem.phi_,V)
        phi = TrialFunction(V)
        varphi = TestFunction(V)

        t = problem.t           
        dt = Constant(SolverBase._get_time_step(self,problem,VV,t))
        T = problem.T

        self.update(problem,phi_,t,float(dt),0)  # store the solution at time 0
        
        v = SolverBase._get_velocity(self,problem,VV,t) # velocity field
        (ib,ib_value,bParts) = problem.get_inflowBCS()  # ib_value has bc conditions on the inflow part
        g = Function(V)                                 # inflow bcs for the weak form
        n = FacetNormal(mesh)                           
     
        bcStrong = self.bcStrong != 0
        if self.type == 0: #bp w
            a = phi*varphi*dx \
                -(1./6.)*dt**2*(-inner(v,nabla_grad(phi))*inner(v,nabla_grad(varphi))*dx + inner(v,n)*inner(v,nabla_grad(phi))*varphi*ds)

            l = phi_*varphi*dx \
                -dt*(inner(v,nabla_grad(phi_))*varphi*dx + 0.5*(abs(inner(v,n))-inner(v,n))*(phi_-g)*varphi*ds)\
                +(1./3.)*dt**2*(-inner(v,nabla_grad(phi_))*inner(v,nabla_grad(varphi))*dx + inner(v,n)*inner(v,nabla_grad(phi_))*varphi*ds)

        elif self.type == 1: #nbp
            a = phi*varphi*dx \
                -(1./6.)*dt**2*inner(v,nabla_grad(inner(v,nabla_grad(phi))))*varphi*dx

            l = phi_*varphi*dx \
                -dt*inner(v,nabla_grad(phi_))*varphi*dx\
                +(1./3.)*dt**2*inner(v,nabla_grad(inner(v,nabla_grad(phi_))))*varphi*dx

        elif self.type == 2: #dbp
            a = phi*varphi*dx \
                -(1./6.)*dt**2*(-inner(v,nabla_grad(phi))*inner(v,nabla_grad(varphi))*dx + inner(v,n)*inner(v,nabla_grad(phi))*varphi*ds)

            l = phi_*varphi*dx \
                -dt*inner(v,nabla_grad(phi_))*varphi*dx\
                +(1./3.)*dt**2*(-inner(v,nabla_grad(phi_))*inner(v,nabla_grad(varphi))*dx + inner(v,n)*inner(v,nabla_grad(phi_))*varphi*ds)

        elif self.type == 3: #bp + no surface term
             a = phi*varphi*dx \
                -(1./6.)*dt**2*(-inner(v,nabla_grad(phi))*inner(v,nabla_grad(varphi))*dx)

             l = phi_*varphi*dx \
                 -dt*inner(v,nabla_grad(phi_))*varphi*dx\
                +(1./3.)*dt**2*(-inner(v,nabla_grad(phi_))*inner(v,nabla_grad(varphi))*dx)

        A = assemble(a)     
        phi = Function(V)
   
        t = 0
        count = 0

        while t <= T:
            t += float(dt)
            count += 1
            g.vector()[:] = interpolate(ib_value(t),V).vector().array() # update inflow values
            
            L = assemble(l) 
            if bcStrong:
                ibc = DirichletBC(V,ib_value(t),ib)  # apply the inflow boundary conditions 
                ibc.apply(A,L)         
 
            solve(A,phi.vector(),L)
            phi_.assign(phi)
            
            self.update(problem,phi_,t,float(dt),count)

        problem.save_data(self.saveDir,str(self.CFL))

